﻿#include "AudioUtils.h"
#include <vector>
#include <string>
#include <thread>
#include <atomic>
#include <chrono>
#include <memory>
#include <wx/wx.h>
#include <wx/file.h>
#include <wx/stdpaths.h>
#include <wx/filename.h>
#include "atlstr.h"
#include <boost/algorithm/string.hpp>
#include <boost/circular_buffer.hpp>

using hrclock = std::chrono::high_resolution_clock;

namespace util::audio {

class AudioRecorderRMS : public AudioStreamReceiver {
    struct RmsHistory {
        double rms;
    };

public:
    explicit AudioRecorderRMS(AudioRmsReceiver* recv) 
        : rmsReceiver(recv) 
        , rmsHistory(60)
    {}

    void Prepare(double sampleRate, double recordStep) override {}

    void Process(const std::vector<int16_t> &pcms) override {
        if (pcms.size() < 10) {
            return;
        }

        auto tpPcm = hrclock::now();
        ++cntPcmPackage;
        if (cntPcmPackage == 1) {
            tpFirstPcmData = tpPcm;
        }

        int64_t sqrsum = 0;
        for (auto pcm : pcms) {
            sqrsum += static_cast<int32_t>(pcm) * pcm;
        }
        auto average = static_cast<double>(sqrsum) / pcms.size();
        auto rms = TimeSmoothing(std::sqrt(average), lastRMS);
        rmsHistory.push_back(RmsHistory{ rms });
        rms = (rms / 32767.0) * 100;
        if (rmsReceiver) {
            auto timeMs = std::chrono::duration_cast<std::chrono::milliseconds>(tpPcm - tpFirstPcmData);
            //rmsReceiver->OnReceiveRMS(0, static_cast<double>(timeMs.count()) / 1000, rms);
            auto ave = GetAverageRMS();
            if (ave > 0) {
                //rmsReceiver->OnReceiveRMS(1, static_cast<double>(timeMs.count()) / 1000, ave1);
                if (ave > rms) {
                    rms *= 0.5;
                }
                else {
                    rms *= 1.2;
                }
                rms = TimeSmoothing(rms, lastOutputRMS);
            }
            rmsReceiver->OnReceiveRMS(2, static_cast<double>(timeMs.count()) / 1000, rms);
        }
    }

    void Finish() override {}

private:
    double TimeSmoothing(double rms, double& lastRmsVal) {
        constexpr double smoothDown = 0.08, smoothUp = 0.6;
        if (lastRmsVal < rms) {
            rms = lastRmsVal * smoothDown + rms * (1 - smoothDown);
        }
        else {
            rms = lastRmsVal * smoothUp + rms * (1 - smoothUp);
        }
        lastRmsVal = rms;
        return rms;
    }

    double GetAverageRMS(bool percentMode = true) const {
        if (!rmsHistory.full()) {
            return 0.0;
        }
        auto sum = 0.0;
        for (auto& his : rmsHistory) {
            sum += his.rms;
        }
        auto ave = sum / rmsHistory.capacity();
        if (percentMode) {
            ave = (ave / INT16_MAX) * 100;
        }
        return ave;
    }

private:
    double lastRMS = 0, lastOutputRMS = 0;
    boost::circular_buffer<RmsHistory> rmsHistory;

    std::size_t cntPcmPackage = 0;
    hrclock::time_point tpFirstPcmData{};
    AudioRmsReceiver* rmsReceiver = nullptr;
    
};

static std::atomic<bool> stopStreamWorker = false;
static std::thread streamWorkerThread{};

static bool SreamWorker(PaStream* stream, double sampleRate, double recordStep, AudioRmsReceiver* recv) {
    CStringW log;
    std::vector<int16_t> buffer;
    buffer.reserve(static_cast<std::size_t>(recordStep * sampleRate) * 10);
    auto receiver = std::make_unique<AudioRecorderRMS>(recv);
    receiver->Prepare(sampleRate, recordStep);
    Pa_StartStream(stream);

    std::size_t frameCnt = 0;
    hrclock::time_point tpLoopStart = hrclock::now();
    hrclock::time_point tpStart{}, tpEnd{}, tpStreamReadStart{}, tpRecvStart{}, tpRecvEnd{};
    std::chrono::microseconds stepTimeUs(static_cast<int>(recordStep * 1000000));
    while (!stopStreamWorker) {
        auto processTimeUs = std::chrono::duration_cast<std::chrono::microseconds>(tpEnd - tpStart);
        auto streamQueryTimeUs = std::chrono::duration_cast<std::chrono::microseconds>(tpStreamReadStart - tpStart);
        auto streamReadTimeUs = std::chrono::duration_cast<std::chrono::microseconds>(tpRecvStart - tpStreamReadStart);
        auto recvTimeUs = std::chrono::duration_cast<std::chrono::microseconds>(tpRecvEnd - tpRecvStart);
        if (processTimeUs > stepTimeUs) {
            log.Format(
                L"process timeout(ms): A[%.3lf], Q[%.3lf], R[%.3lf], P[%.3lf], Cnt[%d]", 
                static_cast<double>(processTimeUs.count()) / 1000,
                static_cast<double>(streamQueryTimeUs.count()) / 1000,
                static_cast<double>(streamReadTimeUs.count()) / 1000,
                static_cast<double>(recvTimeUs.count()) / 1000,
                frameCnt
            );
            OutputDebugStringW(log.GetString());
        }
        else {
            std::this_thread::sleep_for(stepTimeUs - processTimeUs);
        }
        if (stopStreamWorker) {
            OutputDebugStringW(L"exit signal");
            break;
        }
        
        tpStart = hrclock::now();
        tpStreamReadStart = {};
        tpRecvStart = {};
        tpRecvEnd = {};
        // read data
        frameCnt = Pa_GetStreamReadAvailable(stream);
        if (frameCnt > 0) {
            buffer.resize(frameCnt);
            tpStreamReadStart = hrclock::now();
            auto errorCode = Pa_ReadStream(stream, &buffer[0], static_cast<unsigned long>(buffer.size()));
            if (errorCode != paNoError) {
                log.Format(L"Read error: %s", CA2W(Pa_GetLastHostErrorInfo()->errorText, CP_UTF8).m_psz);
                OutputDebugStringW(log.GetString());
                break;
            }
            // process data
            tpRecvStart = hrclock::now();
            receiver->Process(buffer);
            buffer.clear();
            tpRecvEnd = hrclock::now();
        }
        tpEnd = hrclock::now();
    }
    
    Pa_StopStream(stream);
    receiver->Finish();
    return true;
}

PaDeviceIndex GetLoopBackDevice4DefaultOutput() {
    auto sum = Pa_GetDeviceCount();
    if (sum < 0) {
        return paNoDevice;
    }

    const auto* targetOutDevInfo = Pa_GetDeviceInfo(Pa_GetDefaultOutputDevice());
    std::wstring targetName = CA2W(targetOutDevInfo->name, CP_UTF8).m_psz;
    for (auto i = 0; i < sum; i++) {
        auto devInfo = Pa_GetDeviceInfo(i);
        const auto* hostInfo = Pa_GetHostApiInfo(devInfo->hostApi);
        if (hostInfo && hostInfo->type == paWASAPI && PaWasapi_IsLoopback(i)) {
            std::wstring name = CA2W(devInfo->name, CP_UTF8).m_psz;
            if (boost::istarts_with(name, targetName)) {
                return i;
            }
        }
    }
    return paNoDevice;
}

bool IsAudioRecording() {
    return streamWorkerThread.joinable();
}

bool StopAudioRecording() {
    if (!IsAudioRecording()) {
        return false;
    }
    stopStreamWorker = true;
    streamWorkerThread.join();
    stopStreamWorker = false;
    return true;
}

bool StartAudioRecording(AudioRmsReceiver* recv) {
    const auto devIdx = GetLoopBackDevice4DefaultOutput();
    if (devIdx < 0 || devIdx >= Pa_GetDeviceCount()) {
        return false;
    }
    const auto* devInfo = Pa_GetDeviceInfo(devIdx);
    if (!devInfo) {
        return false;
    }

    if (!recv) {
        recv = dynamic_cast<AudioRmsReceiver*>(wxTheApp->GetTopWindow());
    }

    PaStreamParameters inputParam{};
    inputParam.device = devIdx;
    inputParam.channelCount = 1;
    inputParam.sampleFormat = paInt16;
    inputParam.suggestedLatency = devInfo->defaultLowInputLatency;
    inputParam.hostApiSpecificStreamInfo = nullptr;

    // launch stream worker 
    std::thread thread(
        [](PaStreamParameters input, double sampleRate, double recordStep, AudioRmsReceiver* recv) {
            OutputDebugStringW(L"==== stream worker enter ====");
            PaStream* stream = nullptr;
            PaError errorCode = Pa_OpenStream(
                &stream,
                &input,
                nullptr,
                sampleRate,
                static_cast<unsigned long>(recordStep * sampleRate),
                paClipOff,
                nullptr,
                nullptr
            );
            if (errorCode != paNoError) {
                return;
            }
            SreamWorker(stream, sampleRate, recordStep, recv);
            Pa_CloseStream(stream);
            OutputDebugStringW(L"==== stream worker exit ====");
        },
        inputParam, devInfo->defaultSampleRate, 0.02, recv
    );
    if (!thread.joinable()) {
        return false;
    }
    streamWorkerThread.swap(thread);
    return true;
}
}